package com.zhiziai.aiproxy.service;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Preconditions;
import com.zhiziai.aiproxy.config.TheadPoolConfig;
import com.zhiziai.aiproxy.constants.MidJourneyTaskStatus;
import com.zhiziai.aiproxy.constants.UploadType;
import com.zhiziai.aiproxy.dto.MidJourneyImgDto;
import com.zhiziai.aiproxy.dto.MidJourneyTask;
import com.zhiziai.aiproxy.utils.SignUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

/**
 * @author: liuxinyu06
 * @description:
 * @create: 2022-04-07
 **/
@Slf4j
@Service
public class MidjourneyServiceImpl {

    private String IMAGINEGEN_URL = "http://127.0.0.1:8888/mj/submit/imagine";
    private String IMAGINEGEN_UV_URL = "http://127.0.0.1:8888/mj/submit/simple-change";
    private String TASK_SEARCH_URL = "http://127.0.0.1:8888/mj/task/{id}/fetch";
    public static Map<String,CountDownLatch> lockMap = new HashMap<>();
    public static Map<String, MidJourneyTask> listenMidTaskMap = new HashMap<>();

    @Autowired
    private UploadService uploadService;
    @Autowired
    private ThreadPoolTaskScheduler threadPoolTaskScheduler;

    public List<MidJourneyImgDto> encryptImgGen(Long userId, String prompt, Integer genNum, String sign) throws Exception{
        Preconditions.checkArgument(userId != null && StringUtils.isNotEmpty(prompt) && StringUtils.isNotEmpty(sign), "参数错误");
        Preconditions.checkArgument(genNum != null && genNum > 0 && genNum <=4, "genNum error");
        String calSign = SignUtil.buildSign(userId+"",prompt,genNum+"");
        Preconditions.checkArgument(StringUtils.isNotEmpty(calSign) && calSign.equals(sign),"auth error");
        return imagineGen(userId,prompt,genNum);
    }

    public List<MidJourneyImgDto> encryptImgVariation(Long userId, String taskId,int index,String sign) throws Exception{
        Preconditions.checkArgument(StringUtils.isNotEmpty(taskId) && StringUtils.isNotEmpty(sign), "参数错误");
        String calSign = SignUtil.buildSign(userId+"",taskId,index+"");
        Preconditions.checkArgument(StringUtils.isNotEmpty(calSign) && calSign.equals(sign),"auth error");
        return imagineVariation(userId,taskId,index);
    }

    public List<MidJourneyImgDto> imagineGen(Long userId, String prompt, int genNum) {
        Map<String,Object> query = new HashMap<>();
        query.put("action","IMAGINE");
        query.put("prompt",prompt);
        query.put("state",0);
        String res = HttpRequest.post(IMAGINEGEN_URL)
                .body(JSONObject.toJSONString(query))
                .execute().body();
        log.info("imagineGen res: {}",res);
        Preconditions.checkNotNull(res,"imagineGen post error ");
        JSONObject jo = JSONObject.parseObject(res);
        Integer code = jo.getInteger("code");
        Preconditions.checkArgument(code != null && (code.equals(1) || code.equals(2)),jo.getString("description"));
        List<MidJourneyImgDto> result = new LinkedList<>();
        String taskId = jo.getString("result");
        MidJourneyImgDto midJourneyImgDto = waitCreateImg(taskId,0, "u");
        if(midJourneyImgDto != null){
            result.add(midJourneyImgDto);
            result.addAll(asyGenUvImgAnduploadCos(taskId,genNum,"u",userId));
        }
        return result;
    }

    private List<MidJourneyImgDto> imagineVariation(Long userId,String taskId, int index) {
        List<MidJourneyImgDto> result = new ArrayList<>();
        MidJourneyImgDto imgDto = genUVImg(taskId,index,"v");
        if(imgDto != null){
            return asyGenUvImgAnduploadCos(imgDto.getTaskId(),4,"u",userId);
        }
        return result;
    }

    public MidJourneyImgDto genUVImg(String taskId,int index,String action){
        Map<String,Object> query = new HashMap<>();
        query.put("content",taskId + " "+action+index);
        query.put("state",index);
        String res = HttpRequest.post(IMAGINEGEN_UV_URL)
                .body(JSONObject.toJSONString(query))
                .execute().body();
        log.info("taskId {} , genUImg res: {}",taskId, res);
        if(StringUtils.isNotEmpty(res)){
            JSONObject jo = JSONObject.parseObject(res);
            Integer code = jo.getInteger("code");
            if(code != null && (code.equals(1) || code.equals(2))){
                String newTaskId = jo.getString("result");
                return waitCreateImg(newTaskId,index,action);
            }
        }
        return null;
    }

    private MidJourneyImgDto waitCreateImg(String taskId, int index, String action) {
        MidJourneyImgDto imgDto = null;
        CountDownLatch cdl = new CountDownLatch(1);
        lockMap.put(taskId+"_"+index,cdl);
        log.info("waitCreateImg lock taskId {} , index {} , action {}",taskId,index,action);
        listenTaskTimeoutStatus(taskId,index);
        try{
            cdl.await();
            MidJourneyTask midJourneyTask = listenMidTaskMap.get(taskId+"_"+index);
            if(midJourneyTask != null && midJourneyTask.getStatus().equals(MidJourneyTaskStatus.SUCCESS)){
                imgDto = MidJourneyImgDto.builder().taskId(taskId).imgUrl(midJourneyTask.getImageUrl())
                        .t(action +index).build();
                lockMap.remove(taskId+"_"+index);
                listenMidTaskMap.remove(taskId+"_"+index);
            }
        }catch (Exception e){
            log.error("{} , waitCreateImg error", taskId ,e);
        }
        return imgDto;
    }

    private void listenTaskTimeoutStatus(String taskId,int index){
        threadPoolTaskScheduler.schedule(() -> {
            try{
                log.info("listenTaskTimeoutStatus taskId {} , index {} " ,taskId,index);
                String res = HttpUtil.get(TASK_SEARCH_URL.replace("{id}",taskId));
                log.info("listenTaskTimeoutStatus task {}" ,res);
                if(StringUtils.isNotEmpty(res)){
                    MidJourneyTask midJourneyTask = JSONObject.parseObject(res,MidJourneyTask.class);
                    if(midJourneyTask.getStatus().equals(MidJourneyTaskStatus.NOT_START) ||
                            midJourneyTask.getStatus().equals(MidJourneyTaskStatus.SUBMITTED)
                            || midJourneyTask.getProgress().equals("0%")){
                        MidjourneyServiceImpl.lockMap.get(taskId+"_"+index).countDown();
                    }
                }
            }catch (Exception e){
               log.error("listenTaskTimeoutStatus error " ,e );
            }
        },new Date(System.currentTimeMillis() + 180 * 1000));
    }

    private List<MidJourneyImgDto> asyGenUvImgAnduploadCos(String taskId, int genNum, String uv,Long userId){
        List<CompletableFuture<MidJourneyImgDto>> futures = new ArrayList<>();
        for(int i=1;i<=genNum;i++){
            int finalI = i;
            CompletableFuture<MidJourneyImgDto> cf = CompletableFuture.supplyAsync(() -> {
                try {
                    MidJourneyImgDto midJourneyImgDto = genUVImg(taskId, finalI,uv);
                    if(midJourneyImgDto != null && StringUtils.isNotEmpty(midJourneyImgDto.getImgUrl())){
                        String cosUrl = uploadService.upload(userId,midJourneyImgDto.getImgUrl(), UploadType.MIDJOURNEY.getType());
                        midJourneyImgDto.setImgUrl(cosUrl);
                        log.info("userId {} , taskId {} , cosUrl:{}", userId, taskId, cosUrl);
                        return midJourneyImgDto;
                    }
                } catch (Exception e) {
                    log.error("genUVImgOrupload img error", e);
                }
                return null;
            }, TheadPoolConfig.text2ImgThreadPoll);
            futures.add(cf);
        }
        return futures.stream().map(future -> {
            try {
                return future.get();
            } catch (Exception e) {
                future.cancel(true);
            }
            return null;
        }).filter(Objects::nonNull).collect(Collectors.toList());
    }


    public static void main(String[] args) {
        CountDownLatch cdl = new CountDownLatch(1);
//        ReentrantLock reentrantLock = new ReentrantLock();
        CountDownLatch finalCdl = cdl;
        new Thread(()->{
            try {
                Thread.sleep(10000);
                finalCdl.countDown();
//                reentrantLock.notify();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }).start();
        do {
            try {
                cdl.await();
//                reentrantLock.wait();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            System.out.println("ing....");
            cdl = new CountDownLatch(1);

        }while (true);


    }


}
